from neo4j import GraphDatabase
from neo4j import Session


class Neo4jConfig:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        self.driver.close()

    def get_driver(self):
        return self.driver

# 创建数据库连接配置
neo4j_config = Neo4jConfig(
    uri="neo4j://localhost:7687",
    user="neo4j",
    password="123"
)

class Neo4jDAO:
    def __init__(self):
        self.driver = neo4j_config.get_driver()

    def query(self, cypher_query, parameters=None):
        with self.driver.session() as session:
            result = session.run(cypher_query, parameters)
            return [record.data() for record in result]

# 查询函数
def get_related_nodes(input_data: str):
    query = """
    MATCH (d:Disease {name: $name})-[:HAS_SYMPTOM]->(s:Symptom)
    OPTIONAL MATCH (d)-[:RELATED_TO]->(dept:Department)
    RETURN d.name AS disease, COLLECT(DISTINCT s.name) AS symptoms, COLLECT(DISTINCT dept.name) AS departments
    """
    with Neo4jDAO.session() as session:
        result = session.run(query, name=input_data)
        return result.data()

if __name__ == '__main__':
    person_data = get_related_nodes('高血压')
    print("知识图谱返回的结果: ", person_data)